package com.kele.encryption.advice;

import com.kele.encryption.config.EncryptionProperties;
import com.kele.encryption.constant.EncryptionConstant;
import com.kele.encryption.util.AESUtil;
import com.kele.encryption.util.RSAUtil;
import org.apache.commons.io.IOUtils;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdvice;

import java.io.IOException;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;

@ControllerAdvice
@ConditionalOnProperty(value = "data-encryptor.enable", havingValue = "true")
public class DecryptRequestBodyAdvice implements RequestBodyAdvice {

    private final EncryptionProperties encryptionProperties;

    public DecryptRequestBodyAdvice(EncryptionProperties encryptionProperties) {
        this.encryptionProperties = encryptionProperties;
    }

    @Override
    public boolean supports(MethodParameter methodParameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        // 是否开启加解密  且  方法参数是否由 @RequestBody 注解修饰
        return encryptionProperties.isEnable() && methodParameter.hasParameterAnnotation(RequestBody.class);
    }

    @Override
    public HttpInputMessage beforeBodyRead(HttpInputMessage inputMessage, MethodParameter parameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) throws IOException {
        String encryptedAESKey = inputMessage.getHeaders().getFirst(EncryptionConstant.HTTP_HEADER);
        try {
            String AESKey = RSAUtil.decrypt(encryptedAESKey, encryptionProperties.getPrivateKey());
            String encryptedData = IOUtils.toString(inputMessage.getBody(), StandardCharsets.UTF_8);
            String decryptedData = AESUtil.decrypt(encryptedData, AESKey);
            return new HttpInputMessage() {
                @Override
                public java.io.InputStream getBody() throws IOException {
                    return IOUtils.toInputStream(decryptedData, StandardCharsets.UTF_8);
                }
                @Override
                public org.springframework.http.HttpHeaders getHeaders() {
                    return inputMessage.getHeaders();
                }
            };
        } catch (Exception e) {
            throw new RuntimeException("Decryption failed.", e);
        }
    }

    @Override
    public Object afterBodyRead(Object body, HttpInputMessage inputMessage, MethodParameter parameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        return body;
    }

    @Override
    public Object handleEmptyBody(Object body, HttpInputMessage inputMessage, MethodParameter parameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        return body;
    }
}
